
#include <iostream>
#include <vector>

#include "llvm/ADT/StringRef.h"
#include "llvm/Bitcode/BitcodeWriter.h"
#include "llvm/ExecutionEngine/ExecutionEngine.h"
#include "llvm/ExecutionEngine/GenericValue.h"
#include "llvm/ExecutionEngine/Orc/LLJIT.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Verifier.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/DynamicLibrary.h"

/**
 * llvm调用静态库里的函数
 */

bool initLLVM()
{
	llvm::InitializeNativeTarget();
	llvm::InitializeNativeTargetAsmPrinter();
	llvm::InitializeNativeTargetAsmParser();

	return true;
}

llvm::Function* wrapCallExternFunc(llvm::Module* module)
{
	llvm::LLVMContext& context = module->getContext();

	// 定义函数地址
	llvm::sys::DynamicLibrary::LoadLibraryPermanently("./libdlib.so");
	//llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);

	// 获取函数地址
	void* function_address = llvm::sys::DynamicLibrary::SearchForAddressOfSymbol("func_sub");

	assert(function_address && "function address not found!");

	// 创建函数类型
	llvm::FunctionType* FT = llvm::FunctionType::get(llvm::Type::getInt32Ty(context), {
		llvm::Type::getInt32Ty(context),
		llvm::Type::getInt32Ty(context)
	},
	false);

	// 创建函数指针
	llvm::Constant* function_pointer = llvm::ConstantExpr::getBitCast(
	                                           llvm::ConstantInt::get(llvm::Type::getInt64Ty(context), reinterpret_cast<uint64_t>(function_address)),
	                                           llvm::PointerType::getUnqual(FT));

	/**
	 * 创建一个call_func函数，调用func_sub，如下：
	 int func_call(int a, int b) {
	    return func_sub(a, b);
	 }
	 */
	std::vector<llvm::Type*> Integers(2, llvm::Type::getInt32Ty(context));
	llvm::FunctionType* FT_caller = llvm::FunctionType::get(llvm::Type::getInt32Ty(context), Integers, false);
	llvm::Function* func_call = llvm::Function::Create(FT_caller, llvm::Function::ExternalLinkage, "func_call", module);

	// 创建 LLVM 基本块
	llvm::BasicBlock* BB = llvm::BasicBlock::Create(context, "entry", func_call);

	// 生成 LLVM 代码
	llvm::IRBuilder<> builder(BB);

	// 调用函数
	std::vector<llvm::Value*> args;
	auto arg = func_call->arg_begin();
	args.push_back(&(*arg));
	arg = std::next(arg);
	args.push_back(&(*arg));

	llvm::Value* call = builder.CreateCall(FT, function_pointer, args);
	builder.CreateRet(call);

	llvm::verifyFunction(*func_call);

	return func_call;
}

int main(int argc, char* argv[])
{
	if (argc != 3) {
		std::cout << "Usage: " << argv[0] << " <a> <b>" << std::endl;
		return 0;
	}

	int a = std::atoi(argv[1]);
	int b = std::atoi(argv[2]);

	initLLVM();

	// 创建一个LLJIT实例
	auto JIT = llvm::cantFail(llvm::orc::LLJITBuilder().create());

	// 创建一个新的LLVM模块
	llvm::orc::ThreadSafeContext ctx(std::make_unique<llvm::LLVMContext>());
	llvm::LLVMContext* context = ctx.getContext();
	auto module = std::make_unique<llvm::Module>("test", *context);
	module->setDataLayout(JIT->getDataLayout());

	wrapCallExternFunc(module.get());

	module->dump();

	// 保存bitcode文件
	//std::error_code EC;
	//llvm::raw_fd_ostream output("output.bc", EC);
	//std::cout << EC << std::endl;
	//llvm::WriteBitcodeToFile(*module.get(), output);
	//output.flush();

	auto tsm = llvm::orc::ThreadSafeModule(std::move(module), ctx);
	cantFail(JIT->addIRModule(std::move(tsm)));

	// 调用call_func
	auto sym = JIT->lookup("func_call");
	assert(sym && "Function not found");

	int (*FP)(int, int) = (int (*)(int, int))(sym->getAddress());
	int result = FP(a, b);

	std::cout << "result: a - b = " << result << std::endl;

	return 0;
}

